import numpy as np
import airsim
class TrajectoryGenerator:
    def __init__(self, total_time):
        self.total_time = total_time

    def get_desired_state(self, t):
        # 参数设置
        radius = 6.0  # 水平螺旋半径（米）
        omega = 0.2  # 水平旋转角速度（rad/s）
        z_velocity = 0.5  # 垂直上升速度（米/秒）

        # 生成螺旋轨迹
        x_d = radius * np.cos(omega * t)
        y_d = radius * np.sin(omega * t)
        z_d = z_velocity * t  # 垂直方向线性上升

        # 计算速度
        vx_d = -radius * omega * np.sin(omega * t)
        vy_d = radius * omega * np.cos(omega * t)
        vz_d = z_velocity  # 垂直速度恒定

        # 计算角度和角速度
        roll_d = 0.0  # 滚转角（保持水平）
        pitch_d = 0.0  # 俯仰角（保持水平）
        yaw_d = omega * t + np.pi / 2  # 航向角指向切线方向

        # 角速度
        p_d = 0.0  # 滚转角速度
        q_d = 0.0  # 俯仰角速度
        r_d = omega  # 偏航角速度（与水平旋转角速度一致）

        # 合并所有状态
        return np.array([x_d, y_d, z_d,
                         vx_d, vy_d, vz_d,
                         roll_d, pitch_d, yaw_d,
                         p_d, q_d, r_d])
if __name__ == '__main__':
    total_time = 40
    quadrotor=0.2
    traj_gen = TrajectoryGenerator(total_time=total_time)
    n_states=12
    N=8
    p_list = []

    num_steps = int(total_time / quadrotor)
    for i in range(num_steps):
        t = i * quadrotor
        desired_states=traj_gen.get_desired_state(t)
        p_list.append(airsim.Vector3r(desired_states[0],desired_states[1],-desired_states[2]))
    client = airsim.MultirotorClient()
    client.enableApiControl(True)
    client.takeoffAsync().join()
    client.simPlotLineStrip(p_list, color_rgba=[0, 1, 0, 0], is_persistent=True)
    client.simSetTraceLine(color_rgba=[1, 0, 0], thickness=5)
    state = client.getMultirotorState()